{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e038be7e",
   "metadata": {},
   "source": [
    "## 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "d1d4a9e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "def _simplify_add_sub(expr_str):\n",
    "    \"\"\"精确化简括号内的加减法，保留所有括号结构\"\"\"\n",
    "    import re\n",
    "    \n",
    "    def simplify_inner(inner):\n",
    "        # 提取所有带符号的数字\n",
    "        numbers = re.findall(r'([+-]\\s*\\d+)', inner)\n",
    "        # 计算总和\n",
    "        total = sum(int(num.replace(' ', '')) for num in numbers)\n",
    "        # 移除所有数字项\n",
    "        simplified = re.sub(r'[+-]\\s*\\d+', '', inner).strip()\n",
    "        # 清理多余操作符\n",
    "        simplified = re.sub(r'^\\s*[+-]\\s*', '', simplified)\n",
    "        simplified = re.sub(r'\\s*[+-]\\s*$', '', simplified)\n",
    "        # 处理结果\n",
    "        if not simplified:  # 只有数字的情况\n",
    "            return str(total)\n",
    "        if total > 0:\n",
    "            return f\"{simplified} + {total}\" if simplified else f\"{total}\"\n",
    "        elif total < 0:\n",
    "            return f\"{simplified} - {abs(total)}\" if simplified else f\"{total}\"\n",
    "        else:\n",
    "            return simplified if simplified else \"0\"\n",
    "    \n",
    "    # 处理所有括号内的内容\n",
    "    result = expr_str\n",
    "    while True:\n",
    "        new_result = re.sub(r'\\(([^()]+)\\)', \n",
    "                          lambda m: f\"({simplify_inner(m.group(1))})\", \n",
    "                          result)\n",
    "        if new_result == result:\n",
    "            break\n",
    "        result = new_result\n",
    "    \n",
    "    # 清理多余空格\n",
    "    result = re.sub(r'\\s+', ' ', result).strip()\n",
    "    return result"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea196039",
   "metadata": {},
   "source": [
    "## 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca4e29ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ast\n",
    "import operator\n",
    "\n",
    "def _simplify_add_sub(expr_str):\n",
    "    tree = ast.parse(expr_str, mode='eval')\n",
    "\n",
    "    def simplify(node):\n",
    "        if isinstance(node, ast.BinOp) and isinstance(node.op, (ast.Add, ast.Sub)):\n",
    "            left = simplify(node.left)\n",
    "            right = simplify(node.right)\n",
    "\n",
    "            if isinstance(left, ast.Constant) and isinstance(right, ast.Constant):\n",
    "                op_func = operator.add if isinstance(node.op, ast.Add) else operator.sub\n",
    "                return ast.Constant(value=op_func(left.value, right.value))\n",
    "\n",
    "            terms = flatten_add_sub(ast.BinOp(left=left, op=node.op, right=right))\n",
    "            simplified_terms = combine_constants(terms)\n",
    "            return rebuild_ast_from_terms(simplified_terms)\n",
    "\n",
    "        elif isinstance(node, ast.BinOp) and isinstance(node.op, ast.FloorDiv):\n",
    "            return ast.BinOp(left=simplify(node.left), op=node.op, right=simplify(node.right))\n",
    "        elif isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub):\n",
    "            operand = simplify(node.operand)\n",
    "            if isinstance(operand, ast.Constant):\n",
    "                return ast.Constant(value=-operand.value)\n",
    "            return ast.UnaryOp(op=node.op, operand=operand)\n",
    "        elif isinstance(node, ast.Constant) or isinstance(node, ast.Name):\n",
    "            return node\n",
    "        elif isinstance(node, ast.Expression):\n",
    "            return simplify(node.body)\n",
    "        else:\n",
    "            return node\n",
    "\n",
    "    def flatten_add_sub(node):\n",
    "        result = []\n",
    "        def helper(n, sign=1):\n",
    "            if isinstance(n, ast.BinOp):\n",
    "                if isinstance(n.op, ast.Add):\n",
    "                    helper(n.left, sign)\n",
    "                    helper(n.right, sign)\n",
    "                elif isinstance(n.op, ast.Sub):\n",
    "                    helper(n.left, sign)\n",
    "                    helper(n.right, -sign)\n",
    "                else:\n",
    "                    result.append((sign, simplify(n)))\n",
    "            else:\n",
    "                result.append((sign, simplify(n)))\n",
    "        helper(node)\n",
    "        return result\n",
    "\n",
    "    def combine_constants(terms):\n",
    "        const_sum = 0\n",
    "        new_terms = []\n",
    "        for sign, term in terms:\n",
    "            if isinstance(term, ast.Constant):\n",
    "                const_sum += sign * term.value\n",
    "            else:\n",
    "                new_terms.append((sign, term))\n",
    "        if const_sum != 0:\n",
    "            new_terms.append((1, ast.Constant(value=const_sum)))\n",
    "        return new_terms\n",
    "\n",
    "    def rebuild_ast_from_terms(terms):\n",
    "        if not terms:\n",
    "            return ast.Constant(value=0)\n",
    "        node = None\n",
    "        for sign, term in terms:\n",
    "            if sign == -1:\n",
    "                term = ast.UnaryOp(op=ast.USub(), operand=term)\n",
    "            if node is None:\n",
    "                node = term\n",
    "            else:\n",
    "                node = ast.BinOp(left=node, op=ast.Add(), right=term)\n",
    "        return node\n",
    "\n",
    "    def clean_unary_minus(s):\n",
    "        # 把 `+ -` 变成 `-`\n",
    "        return s.replace('+ -', '- ')\n",
    "\n",
    "    simplified_tree = simplify(tree)\n",
    "    result = ast.unparse(simplified_tree)\n",
    "    return clean_unary_minus(result)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5d8aec4",
   "metadata": {},
   "source": [
    "## 3 ok"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "c5cc3c1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ast\n",
    "import operator\n",
    "\n",
    "def _simplify_add_sub(expr_str):\n",
    "    tree = ast.parse(expr_str, mode='eval')\n",
    "\n",
    "    def simplify(node):\n",
    "        if isinstance(node, ast.BinOp):\n",
    "            left = simplify(node.left)\n",
    "            right = simplify(node.right)\n",
    "            \n",
    "            if isinstance(node.op, (ast.Add, ast.Sub)):\n",
    "                if isinstance(left, ast.Constant) and isinstance(right, ast.Constant):\n",
    "                    op_func = operator.add if isinstance(node.op, ast.Add) else operator.sub\n",
    "                    return ast.Constant(value=op_func(left.value, right.value))\n",
    "\n",
    "                terms = flatten_add_sub(ast.BinOp(left=left, op=node.op, right=right))\n",
    "                simplified_terms = combine_constants(terms)\n",
    "                return rebuild_ast_from_terms(simplified_terms)\n",
    "            \n",
    "            elif isinstance(node.op, ast.Mult):\n",
    "                # 如果两边都是常数，直接计算\n",
    "                if isinstance(left, ast.Constant) and isinstance(right, ast.Constant):\n",
    "                    return ast.Constant(value=left.value * right.value)\n",
    "                # 否则保持原样\n",
    "                return ast.BinOp(left=left, op=node.op, right=right)\n",
    "            \n",
    "            elif isinstance(node.op, ast.FloorDiv):\n",
    "                return ast.BinOp(left=left, op=node.op, right=right)\n",
    "            \n",
    "        elif isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub):\n",
    "            operand = simplify(node.operand)\n",
    "            if isinstance(operand, ast.Constant):\n",
    "                return ast.Constant(value=-operand.value)\n",
    "            return ast.UnaryOp(op=node.op, operand=operand)\n",
    "        elif isinstance(node, ast.Constant) or isinstance(node, ast.Name):\n",
    "            return node\n",
    "        elif isinstance(node, ast.Expression):\n",
    "            return simplify(node.body)\n",
    "        else:\n",
    "            return node\n",
    "\n",
    "    def flatten_add_sub(node):\n",
    "        result = []\n",
    "        def helper(n, sign=1):\n",
    "            if isinstance(n, ast.BinOp):\n",
    "                if isinstance(n.op, ast.Add):\n",
    "                    helper(n.left, sign)\n",
    "                    helper(n.right, sign)\n",
    "                elif isinstance(n.op, ast.Sub):\n",
    "                    helper(n.left, sign)\n",
    "                    helper(n.right, -sign)\n",
    "                else:\n",
    "                    result.append((sign, simplify(n)))\n",
    "            else:\n",
    "                result.append((sign, simplify(n)))\n",
    "        helper(node)\n",
    "        return result\n",
    "\n",
    "    def combine_constants(terms):\n",
    "        const_sum = 0\n",
    "        new_terms = []\n",
    "        for sign, term in terms:\n",
    "            if isinstance(term, ast.Constant):\n",
    "                const_sum += sign * term.value\n",
    "            else:\n",
    "                new_terms.append((sign, term))\n",
    "        if const_sum != 0:\n",
    "            new_terms.append((1, ast.Constant(value=const_sum)))\n",
    "        return new_terms\n",
    "\n",
    "    def rebuild_ast_from_terms(terms):\n",
    "        if not terms:\n",
    "            return ast.Constant(value=0)\n",
    "        node = None\n",
    "        for sign, term in terms:\n",
    "            if sign == -1:\n",
    "                term = ast.UnaryOp(op=ast.USub(), operand=term)\n",
    "            if node is None:\n",
    "                node = term\n",
    "            else:\n",
    "                node = ast.BinOp(left=node, op=ast.Add(), right=term)\n",
    "        return node\n",
    "\n",
    "    def clean_unary_minus(s):\n",
    "        # 把 `+ -` 变成 `-`\n",
    "        return s.replace('+ -', '- ')\n",
    "\n",
    "    simplified_tree = simplify(tree)\n",
    "    result = ast.unparse(simplified_tree)\n",
    "    return clean_unary_minus(result)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2000b80",
   "metadata": {},
   "source": [
    "## test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "7bb61b51",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'((ninetoothed_tensor_0_size_0 - 1) // 64 + 1) * ((ninetoothed_tensor_2_size_1 - 1) // 64 + 1)'"
      ]
     },
     "execution_count": 70,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_simplify_add_sub(\"((ninetoothed_tensor_0_size_0 - 63 - 1 + 64 - 1) // 64 + 1 - 1 + 1 - 1 + 1) * ((ninetoothed_tensor_2_size_1 - 63 - 1 + 64 - 1) // 64 + 1)\") "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "id": "bdf9b71f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'(ninetoothed_tensor_0_size_0 - 1) // 64 + 1'"
      ]
     },
     "execution_count": 71,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_simplify_add_sub('(ninetoothed_tensor_0_size_0 - 63 - 1 + 64 - 1) // 64 + 1 - 1 + 1 - 1 + 1') #'(ninetoothed_tensor_0_size_0 - 1) // 64 + 1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "f8351cc5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'x + y + 6'"
      ]
     },
     "execution_count": 72,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_simplify_add_sub('((x + 10 - 5) + (y - 3 + 1)) - 2 + 5') # '((x + 5) + (y - 2)) + 3'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "26eb9cad",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'90'"
      ]
     },
     "execution_count": 73,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_simplify_add_sub('(100 - 50 + 20) + 30 - 10') # '(70) + 20'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "c584ae01",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'x + 6'"
      ]
     },
     "execution_count": 74,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_simplify_add_sub('x + 10 - 5 + 3 - 2') # 'x + 6'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "a643d187",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'(x + 2) // 2 + 3'"
      ]
     },
     "execution_count": 75,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_simplify_add_sub('(x + 5 - 3) // 2 + 4 - 1') #'(x + 2) // 2 + 3'"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "my_env_torch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
